##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
## Created for Potsdam dataset
## Based on DefaultLoader
## Copyright (c) 2025
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import pdb

import numpy as np
from torch.utils import data

from lib.utils.helpers.image_helper import ImageHelper
from lib.extensions.parallel.data_container import DataContainer
from lib.utils.tools.logger import Logger as Log


class PotsdamLoader(data.Dataset):
    def __init__(self, root_dir, aug_transform=None, dataset=None,
                 img_transform=None, label_transform=None, configer=None):
        self.configer = configer
        self.aug_transform = aug_transform
        self.img_transform = img_transform
        self.label_transform = label_transform
        self.img_list, self.label_list, self.name_list = self.__list_dirs(root_dir, dataset)
        size_mode = self.configer.get(dataset, 'data_transformer')['size_mode']
        self.is_stack = size_mode != 'diverse_size'

        Log.info('{} {}'.format(dataset, len(self.img_list)))

    def __len__(self):
        return len(self.img_list)

    def __getitem__(self, index):
        img = ImageHelper.read_image(self.img_list[index],
                                     tool=self.configer.get('data', 'image_tool'),
                                     mode=self.configer.get('data', 'input_mode'))
        img_size = ImageHelper.get_size(img)
        labelmap = ImageHelper.read_image(self.label_list[index],
                                          tool=self.configer.get('data', 'image_tool'), mode='P')
        
        # 移除调试信息
        
        # Potsdam特定的标签编码
        # 注意：预处理后的数据已经是0-5格式，不需要再编码
        if self.configer.exists('data', 'label_list'):
            # 检查是否需要编码（原始数据包含29,76,150,179,226,255）
            original_values = np.unique(ImageHelper.tonp(labelmap))
            if any(val in [29, 76, 150, 179, 226, 255] for val in original_values):
                # 原始数据，需要编码
                labelmap = self._encode_potsdam_label(labelmap)
                pass
            else:
                # 预处理后的数据，已经是0-5格式，直接使用
                pass

        if self.configer.exists('data', 'reduce_zero_label'):
            labelmap = self._reduce_zero_label(labelmap)

        ori_target = ImageHelper.tonp(labelmap)
        # 对于Potsdam数据集，不要将255设置为-1，因为255已经映射到类别5
        # ori_target[ori_target == 255] = -1  # 注释掉这行

        if self.aug_transform is not None:
            img, labelmap = self.aug_transform(img, labelmap=labelmap)
            # 移除调试信息

        border_size = ImageHelper.get_size(img)

        if self.img_transform is not None:
            img = self.img_transform(img)

        if self.label_transform is not None:
            labelmap = self.label_transform(labelmap)
            # 移除调试信息

        meta = dict(
            ori_img_size=img_size,
            border_size=border_size,
            ori_target=ori_target
        )
        return dict(
            img=DataContainer(img, stack=self.is_stack),
            labelmap=DataContainer(labelmap, stack=self.is_stack),
            meta=DataContainer(meta, stack=False, cpu_only=True),
            name=DataContainer(self.name_list[index], stack=False, cpu_only=True),
        )

    def _reduce_zero_label(self, labelmap):
        if not self.configer.get('data', 'reduce_zero_label'):
            return labelmap

        labelmap = np.array(labelmap)
        encoded_labelmap = labelmap - 1
        if self.configer.get('data', 'image_tool') == 'pil':
            encoded_labelmap = ImageHelper.np2img(encoded_labelmap.astype(np.uint8))

        return encoded_labelmap

    def _encode_potsdam_label(self, labelmap):
        """Potsdam特定的标签编码
        将原始标签值 [29, 76, 150, 179, 226, 255] 映射到 [0, 1, 2, 3, 4, 5]
        """
        labelmap = np.array(labelmap)
        shape = labelmap.shape
        encoded_labelmap = np.ones(shape=(shape[0], shape[1]), dtype=np.float32) * 255
        
        # Potsdam标签映射
        label_mapping = {
            29: 0,   # Impervious surfaces
            76: 1,   # Building  
            150: 2,  # Low vegetation
            179: 3,  # Tree
            226: 4,  # Car
            255: 5   # Clutter/background
        }
        
        for original_label, new_label in label_mapping.items():
            mask = labelmap == original_label
            if mask.any():
                encoded_labelmap[mask] = new_label

        if self.configer.get('data', 'image_tool') == 'pil':
            encoded_labelmap = ImageHelper.np2img(encoded_labelmap.astype(np.uint8))

        return encoded_labelmap

    def __list_dirs(self, root_dir, dataset):
        """列出图像和标签文件"""
        img_list = list()
        label_list = list()
        name_list = list()

        # 构建数据路径
        if dataset == 'train':
            img_dir = os.path.join(root_dir, 'processed', 'train', 'images')
            label_dir = os.path.join(root_dir, 'processed', 'train', 'labels')
        elif dataset == 'val':
            img_dir = os.path.join(root_dir, 'processed', 'val', 'images')
            label_dir = os.path.join(root_dir, 'processed', 'val', 'labels')
        elif dataset == 'test':
            img_dir = os.path.join(root_dir, 'processed', 'test', 'images')
            label_dir = os.path.join(root_dir, 'processed', 'test', 'labels')
        else:
            raise ValueError(f"Unknown dataset: {dataset}")

        if not os.path.exists(img_dir):
            raise FileNotFoundError(f"Image directory not found: {img_dir}")
        if not os.path.exists(label_dir):
            raise FileNotFoundError(f"Label directory not found: {label_dir}")

        # 获取所有图像文件
        img_files = [f for f in os.listdir(img_dir) if f.endswith('.png')]
        img_files.sort()

        for img_file in img_files:
            # 构建对应的标签文件名
            label_file = img_file  # 图像和标签文件名相同
            
            img_path = os.path.join(img_dir, img_file)
            label_path = os.path.join(label_dir, label_file)
            
            if os.path.exists(label_path):
                img_list.append(img_path)
                label_list.append(label_path)
                name_list.append(os.path.splitext(img_file)[0])
            else:
                Log.warning(f"Label file not found: {label_path}")

        return img_list, label_list, name_list
